from coin.blockchain.web3.tron.native_private_client import (
    TronNativePrivateClient as nprivc, TronWalletException)
from coin.blockchain.web3.tron.constants import (
    TRON_TOKEN_TO_ADDRESS, TOKEN_TO_UNIT)
from coin.blockchain.request_util import (
    retry_error, parse_currency_balance, parse_coin_key)
from coin.exchange.base.kr_rest.private_client_base import \
    ExWalletPrivateClientBase


class TronPrivateClient(ExWalletPrivateClientBase):
  def __init__(self, *, blockchain_name, node_info=None, api_key=None):
    assert blockchain_name == 'tron'
    self._blockchain_name = blockchain_name
    node_info = node_info or 'https://api.trongrid.io/'
    self._nprivc = nprivc(node_info=node_info, api_key=api_key)
    self._token_to_address = TRON_TOKEN_TO_ADDRESS
    self._token_unit = TOKEN_TO_UNIT

  def _get_trx_balance(self, address, end_height):
    return self._nprivc.get_trx_balance(address, end_height)

  def get_token_balance(self, *, address, token_address, end_height):
    return self._nprivc.get_token_balance(address, token_address, end_height)

  def query_account_balance_impl(self, *, address, currency, end_height=None):
    token_addres = self._token_to_address[currency]
    balance = retry_error(
        self.get_token_balance,
        address=address,
        token_address=token_addres,
        end_height=end_height)
    balance = parse_currency_balance(
      currency, self._token_unit[currency], balance)
    return balance

  def get_latest_block_number_timestamp(self):
    return retry_error(self._nprivc.get_latest_block_number_timestamp)

  def get_block_timestamp(self, block_number):
    return retry_error(self._nprivc.get_block_timestamp, block_number=block_number)


if __name__ == '__main__':
  key = parse_coin_key('../../coin_key/view/ledger02/hw_wallet/view_key.json')
  wallet_address = next(
    chain.access_key for chain in key.chain_keys if chain.blockchain == 'tron')
  tron_client = TronPrivateClient(node_info='https://api.trongrid.io/',
                                  blockchain_name='tron',
                                  api_key='c1673166-4397-4986-be73-5eef10c04caa')
  block_num, block_ts = tron_client.get_latest_block_number_timestamp()
  print(tron_client.query_account_balance_impl(address=wallet_address,
                                               currency='USDT', end_height=block_num))
